{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3IrnqNnWO04A"
   },
   "source": [
    "# Ungraded Lab: Feature Engineering with Weather Data\n",
    "\n",
    "In this 1st exercise on feature engineering with time series data, you will practice data transformation with the [Weather Dataset](https://www.bgc-jena.mpg.de/wetter/) recorded by the [Max Planck Institute for Biogeochemistry](https://www.bgc-jena.mpg.de/). You will be using [tf.Transform](https://www.tensorflow.org/tfx/transform/get_started) here instead of TFX because, as of this version (1.4), it is more straightforward to preserve the sequence of your records using this framework. If you remember, TFX by default always shuffles the data when ingesting through the `ExampleGen` component and that is not ideal when preparing data for forecasting applications.\n",
    "\n",
    "This dataset has 10-minute intervals of 14 different features such as air temperature, atmospheric pressure, and humidity. For this lab, you will use only the data collected between 2009 and 2016. This section of the dataset was prepared by François Chollet for his book *Deep Learning with Python*.\n",
    "\n",
    "The table below shows the column names, their value formats, and their description.\n",
    "\n",
    "Index| Features      |Format             |Description\n",
    "-----|---------------|-------------------|-----------------------\n",
    "1    |Date Time      |01.01.2009 00:10:00|Date-time reference\n",
    "2    |p (mbar)       |996.52             |The pascal SI derived unit of pressure used to quantify internal pressure. Meteorological reports typically state atmospheric pressure in millibars.\n",
    "3    |T (degC)       |-8.02              |Temperature in Celsius\n",
    "4    |Tpot (K)       |265.4              |Temperature in Kelvin\n",
    "5    |Tdew (degC)    |-8.9               |Temperature in Celsius relative to humidity. Dew Point is a measure of the absolute amount of water in the air, the DP is the temperature at which the air cannot hold all the moisture in it and water condenses.\n",
    "6    |rh (%)         |93.3               |Relative Humidity is a measure of how saturated the air is with water vapor, the %RH determines the amount of water contained within collection objects.\n",
    "7    |VPmax (mbar)   |3.33               |Saturation vapor pressure\n",
    "8    |VPact (mbar)   |3.11               |Vapor pressure\n",
    "9    |VPdef (mbar)   |0.22               |Vapor pressure deficit\n",
    "10   |sh (g/kg)      |1.94               |Specific humidity\n",
    "11   |H2OC (mmol/mol)|3.12               |Water vapor concentration\n",
    "12   |rho (g/m ** 3) |1307.75            |Airtight\n",
    "13   |wv (m/s)       |1.03               |Wind speed\n",
    "14   |max. wv (m/s)  |1.75               |Maximum wind speed\n",
    "15   |wd (deg)       |152.3              |Wind direction in degrees\n",
    "\n",
    "You will perform data preprocessing so that the features can be used to train an LSTM using TensorFlow and Keras downstream. You will not be asked to train a model as the focus is on feature preprocessing.\n",
    "\n",
    "Upon completion, you will have\n",
    "\n",
    "* Explored and visualized the weather time series dataset and declared its schema\n",
    "* Transformed the data for modeling using [tf.Transform](https://www.tensorflow.org/tfx/transform/get_started)\n",
    "* Prepared training dataset windows from the output of `tf.Transform`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QDhCGkJoQDEN"
   },
   "source": [
    "## Install Packages\n",
    "\n",
    "First, you will install the required packages for this lab. Most are included in the `tensorflow_transform` package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qONoNgHNWKt3"
   },
   "outputs": [],
   "source": [
    "!pip install tensorflow_transform==1.12.0\n",
    "!pip install apache-beam==2.40.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xvoAuLGMQXhL"
   },
   "source": [
    "*Note: In Google Colab, you need to restart the runtime at this point to finalize updating the packages you just installed. You can do so by clicking the `Restart Runtime` at the end of the output cell above (after installation), or by selecting `Runtime > Restart Runtime` in the Menu bar. **Please do not proceed to the next section without restarting.** You can also ignore the errors about version incompatibility of some of the bundled packages because we won't be using those in this notebook.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iyHZtPotQhL0"
   },
   "source": [
    "## Imports\n",
    "\n",
    "Running the imports below should not show any error. Otherwise, please restart your runtime or re-run the package installation cell above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b5isfzOcWbCl"
   },
   "outputs": [],
   "source": [
    "import apache_beam as beam\n",
    "print('Apache Beam version: {}'.format(beam.__version__))\n",
    "\n",
    "import tensorflow as tf\n",
    "print('Tensorflow version: {}'.format(tf.__version__))\n",
    "\n",
    "import tensorflow_transform as tft\n",
    "from tensorflow_transform import beam as tft_beam\n",
    "from tensorflow_transform.tf_metadata import dataset_metadata\n",
    "from tensorflow_transform.tf_metadata import schema_utils\n",
    "print('TensorFlow Transform version: {}'.format(tft.__version__))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ll_BqMx4QyP5"
   },
   "source": [
    "## Download the Data\n",
    "\n",
    "Next, you will download the data and put it in your workspace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Lu4uJ4RgEfE8"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Directory of the raw data files\n",
    "DATA_DIR = '/content/data/'\n",
    "\n",
    "# Download the dataset\n",
    "!wget -nc https://raw.githubusercontent.com/https-deeplearning-ai/MLEP-public/main/course2/week4-ungraded-lab/data/jena_climate_2009_2016.csv -P {DATA_DIR}\n",
    "\n",
    "# Assign data path to a variable for easy reference\n",
    "INPUT_FILE = os.path.join(DATA_DIR, 'jena_climate_2009_2016.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oKDJaZfEQ1yT"
   },
   "source": [
    "You can now preview the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-ohkBeclndhL"
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# Put dataset in a dataframe\n",
    "df = pd.read_csv(INPUT_FILE, header=0, index_col=0)\n",
    "\n",
    "# Preview the last few rows\n",
    "df.tail()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "poxo5wBt7oI7"
   },
   "source": [
    "As you can see above, an observation is recorded every 10 minutes. This means that, for a single hour, you will have 6 observations. Similarly, a single day will contain 144 (6x24) observations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y7-nvQoPgJ32"
   },
   "source": [
    "##Inspect the Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5yRWsz06NdhF"
   },
   "source": [
    "You can then inspect the data to see if there are any issues that you need to address before feeding it to your model. First, you will generate some descriptive statistics using the [`describe()`](https://www.geeksforgeeks.org/python-pandas-dataframe-describe-method/) method of the dataframe.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EvF4w5rGNUnf"
   },
   "outputs": [],
   "source": [
    "df.describe().transpose()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lmOcUJXReqhg"
   },
   "source": [
    "You can see that the min value for `wv (m/s)` and `max. wv(m/s)` is `-9999.0`. Those are pretty intense winds and based on the other data points, these just look like faulty measurements. This is more pronounced when you visualize the data using the utilities below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0QAfG_lJ__JK"
   },
   "outputs": [],
   "source": [
    "# Define feature keys\n",
    "\n",
    "TIMESTAMP_FEATURES = [\"Date Time\"]\n",
    "NUMERIC_FEATURES = [\n",
    "    \"p (mbar)\",\n",
    "    \"T (degC)\",\n",
    "    \"Tpot (K)\",\n",
    "    \"Tdew (degC)\", \n",
    "    \"rh (%)\", \n",
    "    \"VPmax (mbar)\", \n",
    "    \"VPact (mbar)\", \n",
    "    \"VPdef (mbar)\", \n",
    "    \"sh (g/kg)\",\n",
    "    \"H2OC (mmol/mol)\",\n",
    "    \"rho (g/m**3)\",\n",
    "    \"wv (m/s)\",\n",
    "    \"max. wv (m/s)\",\n",
    "    \"wd (deg)\",\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "egnSszZ6Ia6k"
   },
   "outputs": [],
   "source": [
    "#@ title Visualization Utilities\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Color Palette\n",
    "colors = [\n",
    "    \"blue\",\n",
    "    \"orange\",\n",
    "    \"green\",\n",
    "    \"red\",\n",
    "    \"purple\",\n",
    "    \"brown\",\n",
    "    \"pink\",\n",
    "    \"gray\",\n",
    "    \"olive\",\n",
    "    \"cyan\",\n",
    "]\n",
    "\n",
    "# Plots each column as a time series\n",
    "def visualize_plots(dataset, columns):\n",
    "    features = dataset[columns]\n",
    "    fig, axes = plt.subplots(\n",
    "        nrows=len(columns)//2 + len(columns)%2, ncols=2, figsize=(15, 20), dpi=80, facecolor=\"w\", edgecolor=\"k\"\n",
    "    )\n",
    "    for i, col in enumerate(columns):\n",
    "        c = colors[i % (len(colors))]\n",
    "        t_data = dataset[col]\n",
    "        t_data.index = dataset.index\n",
    "        t_data.head()\n",
    "        ax = t_data.plot(\n",
    "            ax=axes[i // 2, i % 2],\n",
    "            color=c,\n",
    "            title=\"{}\".format(col),\n",
    "            rot=25,\n",
    "        )\n",
    "    ax.legend([col])\n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sLa8UAe3JDZ3"
   },
   "outputs": [],
   "source": [
    "# Visualize the dataset\n",
    "visualize_plots(df, NUMERIC_FEATURES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eHdr5yhmUxVy"
   },
   "source": [
    "As you can see, there's a very big downward spike towards `-9999` for the two features mentioned. You may know of different methods in handling outliers but for simplicity in this exercise, you will just set these to `0`. You can visualize the plots again after making this change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0zP5RvoFU7KB"
   },
   "outputs": [],
   "source": [
    "# Set the wind velocity outliers to 0\n",
    "wv = df['wv (m/s)']\n",
    "bad_wv = wv == -9999.0\n",
    "wv[bad_wv] = 0.0\n",
    "\n",
    "# Set the max wind velocity outliers to 0\n",
    "max_wv = df['max. wv (m/s)']\n",
    "bad_max_wv = max_wv == -9999.0\n",
    "max_wv[bad_max_wv] = 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wXkjakrEVEYn"
   },
   "outputs": [],
   "source": [
    "visualize_plots(df, NUMERIC_FEATURES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W4MXx8JWapGY"
   },
   "source": [
    "Take note that you are just visualizing the Pandas dataframe here. You will do this simple data cleaning again later when Tensorflow Transform consumes the raw CSV file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "730ghIrjbXhu"
   },
   "source": [
    "## Feature Engineering\n",
    "\n",
    "Now you will be doing feature engineering. There are several things to note before doing the transformation:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6VtsR_d2hhNH"
   },
   "source": [
    "\n",
    "\n",
    "### Correlated features\n",
    "\n",
    "You may want to drop redundant features to reduce the complexity of your model. Let's see what features are highly correlated with each other by plotting the correlation matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sejFrBtwhjV_"
   },
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "\n",
    "def show_correlation_heatmap(dataframe):\n",
    "    plt.figure(figsize=(20,20))\n",
    "    cor = dataframe.corr()\n",
    "    sns.heatmap(cor, annot=True, cmap=plt.cm.PuBu)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "J7uReJtJhqyM"
   },
   "outputs": [],
   "source": [
    "show_correlation_heatmap(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gsMMBBm8haJS"
   },
   "source": [
    "You can observe that `Tpot (K)`, `Tdew (degC)`, `VPmax(mbar)`, `Vpact(mbar)`, `VPdef (mbar)`, `sh(g/kg)` and `H2OC` are highly positively correlated to the target `T (degC)`. Likewise, `rho` is highly negatively correlated to the target. \n",
    "\n",
    "In the features that are positively correlated, you can see that `VPmax (mbar)` is highly correlated to some features like `Vpact (mbar)`, `Tdew (degC)` and `Tpot (K)`. Hence, for the sake of this exercise you can drop these features and retain `VPmax (mbar)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "63D0d5nhjbQ5"
   },
   "outputs": [],
   "source": [
    "# Features to filter out\n",
    "FEATURES_TO_REMOVE = [\"Tpot (K)\", \"Tdew (degC)\",\"VPact (mbar)\" , \"H2OC (mmol/mol)\", \"max. wv (m/s)\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_GFwevJxkqyY"
   },
   "source": [
    "### Distribution of Wind Data\n",
    "\n",
    "The last column of the data, `wd (deg)`, gives the wind direction in units of degrees. However, angles in this current format do not make good model inputs. 360° and 0° should be close to each other and wrap around smoothly. Direction shouldn't matter if the wind is not blowing. This will be easier for the model to interpret if you convert the wind direction and velocity columns to a wind vector. Observe how sine and cosine are used to generate wind vector features (`Wx` and `Wy`) in the `preprocessing_fn()` later.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tfiFpkMRNduM"
   },
   "source": [
    "### Date Time Feature\n",
    "\n",
    "Dealing with weather, you can expect patterns depending on when the measurements are made. For example, temperatures are generally colder at night and wind velocities might be higher during typhoon season. Thus, the `Date Time` column is a good input to your model to take daily and yearly periodicity into account. \n",
    "\n",
    "To do this, you will first use the [datetime](https://docs.python.org/3/library/datetime.htm) Python module to convert the current date time string format (i.e. _day.month.year hour:minute:second_) to a timestamp with units in seconds. Then, a simple approach to generate a periodic signal is to again use sine and cosine to convert the timestamp to clear \"Time of day\" (`Day sin`, `Day cos`) and \"Time of year\" (`Year sin`, `Year cos`) signals. You will see these conversions in the `clean_fn()` and `preprocessing_fn()`. \n",
    "\n",
    "You can see the `clean_fn()` utility function below that removes wind velocity outliers and converts the date time string to a [Unix timestamp](https://www.unixtimestamp.com/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_Lv_foXhGpFD"
   },
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "# combine features into one list\n",
    "ordered_columns = TIMESTAMP_FEATURES + NUMERIC_FEATURES\n",
    "\n",
    "# index of the date time string\n",
    "date_time_idx = ordered_columns.index(TIMESTAMP_FEATURES[0])\n",
    "\n",
    "# index of the 'wv (m/s)' feature\n",
    "wv_idx = ordered_columns.index('wv (m/s)')\n",
    "\n",
    "def clean_fn(line):\n",
    "  '''\n",
    "  Converts datetime strings in the CSV to Unix timestamps and removes outliers\n",
    "  the wind velocity column. Used as part of\n",
    "  the transform pipeline.\n",
    "\n",
    "  Args:\n",
    "    line (string) - one row of a CSV file\n",
    "  \n",
    "  Returns:\n",
    "\n",
    "  '''\n",
    "\n",
    "  # Split the CSV string to a list\n",
    "  line_split = line.split(b',')\n",
    "\n",
    "  # Decodes the timestamp string to utf-8\n",
    "  date_time_string = line_split[date_time_idx].decode(\"utf-8\")\n",
    "\n",
    "  # Creates a datetime object from the timestamp string\n",
    "  date_time = datetime.strptime(date_time_string, '%d.%m.%Y %H:%M:%S')\n",
    "\n",
    "  # Generates a timestamp from the object\n",
    "  timestamp = datetime.timestamp(date_time)\n",
    "\n",
    "  # Overwrites the string timestamp in the row with the timestamp in seconds\n",
    "  line_split[date_time_idx] = bytes(str(timestamp), 'utf-8')\n",
    "\n",
    "  # Check if wind velocity is an outlier\n",
    "  if line_split[wv_idx] == b'-9999.0':\n",
    "\n",
    "    # Overwrite with default value of 0\n",
    "    line_split[wv_idx] = b'0.0'\n",
    "\n",
    "  # rejoin the list item into one string\n",
    "  mod_line = b','.join(line_split)\n",
    "\n",
    "  return mod_line"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KroNS27tp-Eu"
   },
   "source": [
    "### Create a `tf.Transform` preprocessing_fn\n",
    "\n",
    "With the considerations above, you can now declare your `preprocessing_fn()`. This will be used by Tensorflow Transform to create a transformation graph that will preprocess model inputs. In a nutshell, your preprocessing function will perform the following steps:\n",
    "\n",
    "1. Perform feature selection by deleting the unwanted features. \n",
    "2. Transform wind direction and velocity columns into a wind vector.\n",
    "3. Convert date in timestamp to a usable signal by using `sin` and `cos` to convert the time to clear \"Time of day\" and \"Time of year\" signals. \n",
    "4. Normalize the float features.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "neDQbbY30Sdu"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import math as m\n",
    "\n",
    "def preprocessing_fn(inputs):\n",
    "  \"\"\"Preprocess input columns into transformed columns.\"\"\"\n",
    "  \n",
    "  outputs = inputs.copy()\n",
    "\n",
    "  # Filter redundant features\n",
    "  for key in FEATURES_TO_REMOVE:\n",
    "    del outputs[key]\n",
    "\n",
    "  # Convert degrees to radians\n",
    "  pi = tf.constant(m.pi)\n",
    "  wd_rad = inputs['wd (deg)'] * pi / 180.0\n",
    "\n",
    "  # Calculate the wind x and y components.\n",
    "  outputs['Wx'] = inputs['wv (m/s)'] * tf.math.cos(wd_rad)\n",
    "  outputs['Wy'] = inputs['wv (m/s)'] * tf.math.sin(wd_rad)\n",
    "\n",
    "  # Delete `wv (m/s)` and `wd (deg)` after getting the wind vector\n",
    "  del outputs['wv (m/s)']\n",
    "  del outputs['wd (deg)']\n",
    "\n",
    "  # Get day and year in seconds\n",
    "  day = tf.cast(24*60*60, tf.float32)\n",
    "  year = tf.cast((365.2425)*day, tf.float32)\n",
    "\n",
    "  # Get timestamp feature\n",
    "  timestamp_s = outputs['Date Time']\n",
    "\n",
    "  # Convert timestamps into periodic signals\n",
    "  outputs['Day sin'] = tf.math.sin(timestamp_s * (2 * pi / day))\n",
    "  outputs['Day cos'] = tf.math.cos(timestamp_s * (2 * pi / day))\n",
    "  outputs['Year sin'] = tf.math.sin(timestamp_s * (2 * pi / year))\n",
    "  outputs['Year cos'] = tf.math.cos(timestamp_s * (2 * pi / year))\n",
    "\n",
    "  # Delete timestamp feature\n",
    "  del outputs['Date Time']\n",
    "\n",
    "  # Declare final list of features\n",
    "  FINAL_FEATURE_LIST =  [\"p (mbar)\",\n",
    "    \"T (degC)\",\n",
    "    \"rh (%)\", \n",
    "    \"VPmax (mbar)\", \n",
    "    \"VPdef (mbar)\", \n",
    "    \"sh (g/kg)\",\n",
    "    \"rho (g/m**3)\",\n",
    "    \"Wx\",\n",
    "    \"Wy\",\n",
    "    \"Day sin\",\n",
    "    'Day cos',\n",
    "    'Year sin',\n",
    "    'Year cos'\n",
    "    ]\n",
    "\n",
    "  # Scale all features\n",
    "  for key in FINAL_FEATURE_LIST:\n",
    "    outputs[key] = tft.scale_to_0_1(outputs[key])\n",
    "\n",
    "  return outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3SSRMC37SkyI"
   },
   "source": [
    "## Transform the data\n",
    "\n",
    "You're almost ready to start transforming the data in an Apache Beam pipeline. Before doing so, you will declare just a few more utility functions and variables."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qmEOKqhno-pC"
   },
   "source": [
    "### Train Test Split\n",
    "\n",
    "First, you will define how your dataset will be split. You will use the first 300,000 observations for training and the rest for testing.\n",
    "\n",
    "You will extract the date time object of the 300,000th observation to use it in partitioning the dataset using [`Beam.Partition()`](https://beam.apache.org/documentation/transforms/python/elementwise/partition/). This method expects a `partition_fn()` that returns an integer indicating the partition number. Since you will only need two (i.e. train and test), you will make the function return `0` when it is part of the train split, and `1` for the test. See how this is implemented in the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6PJzIMFgBl-b"
   },
   "outputs": [],
   "source": [
    "# Number of records to include in the train split\n",
    "TRAIN_SPLIT = 300000\n",
    "\n",
    "# Get date time of the last element in the train split\n",
    "date_time_train_boundary = df.iloc[TRAIN_SPLIT - 1].name\n",
    "\n",
    "# Creates a datetime object from the timestamp string\n",
    "date_time_train_boundary = datetime.strptime(date_time_train_boundary, '%d.%m.%Y %H:%M:%S')\n",
    "\n",
    "# Convert date time string to Unix timestamp in seconds\n",
    "date_time_train_boundary = bytes(str(datetime.timestamp(date_time_train_boundary)), 'utf-8')\n",
    "\n",
    "\n",
    "def partition_fn(line, num_partitions):\n",
    "  '''\n",
    "  Partition function to work with Beam.partition\n",
    "\n",
    "  Args:\n",
    "    line (string) - One record in the CSV file.\n",
    "    num_partition (integer) - Number of partitions. Required argument by Beam. Unused in this function.\n",
    "\n",
    "  Returns:\n",
    "    0 or 1 (integer) - 0 if line timestamp is below the date time boundary, 1 otherwise. \n",
    "  '''\n",
    "\n",
    "  # Split the CSV string to a list\n",
    "  line_split = line.split(b',')\n",
    "\n",
    "  # Get the timestamp of the current line\n",
    "  line_dt = line_split[date_time_idx]\n",
    "\n",
    "  # Check if it is above or below the date time boundary\n",
    "  partition_num = int(line_dt > date_time_train_boundary)\n",
    "\n",
    "  return partition_num"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "NWDCARvqnYpz"
   },
   "source": [
    "### Declare Schema for Cleaned Data\n",
    "\n",
    "Just like in the previous labs with TFX, you will want to declare a schema to make sure that your data input is parsed correctly. You can do that with the cell below. Take note that this will be used later after the data cleaning step. Thus, you can expect that the date time feature is now in seconds and assigned to be a float feature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-NpXb7N_vWkB"
   },
   "outputs": [],
   "source": [
    "# Declare feature spec\n",
    "RAW_DATA_FEATURE_SPEC = dict(\n",
    "    [(name, tf.io.FixedLenFeature([], tf.float32))\n",
    "     for name in TIMESTAMP_FEATURES] +\n",
    "    [(name, tf.io.FixedLenFeature([], tf.float32))\n",
    "     for name in NUMERIC_FEATURES]\n",
    ")\n",
    "\n",
    "# Create schema from feature spec\n",
    "RAW_DATA_SCHEMA = tft.tf_metadata.schema_utils.schema_from_feature_spec(RAW_DATA_FEATURE_SPEC)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "79ezG-xidnAP"
   },
   "source": [
    "### Create the `tf.Transform` pipeline\n",
    "\n",
    "Now you can define the TF Transform pipeline. It will follow these major steps:\n",
    "\n",
    "1. Read in the data using the CSV reader\n",
    "1. Remove outliers and reformat timestamps using the `clean_fn`.\n",
    "1. Partition the dataset into train and test splits using the `beam.Partition` transform.\n",
    "1. Preprocess the data splits using the `preprocessing_fn`.\n",
    "1. Write the result as a `TFRecord` of `Example` protos."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xJ_eE8nntSSs"
   },
   "outputs": [],
   "source": [
    "import shutil\n",
    "from tfx_bsl.coders.example_coder import RecordBatchToExamplesEncoder\n",
    "from tfx_bsl.public import tfxio\n",
    "\n",
    "# Directory names for the TF Transform outputs\n",
    "WORKING_DIR = 'transform_dir'\n",
    "TRANSFORM_TRAIN_FILENAME = 'transform_train'\n",
    "TRANSFORM_TEST_FILENAME = 'transform_test'\n",
    "TRANSFORM_TEMP_DIR = 'tft_temp'\n",
    "\n",
    "\n",
    "# The \"with\" block will create a pipeline, and run that pipeline at the exit\n",
    "#   of the block.\n",
    "def read_and_transform_data(working_dir):\n",
    "  '''\n",
    "  Reads a CSV File and preprocesses the data using TF Transform\n",
    "\n",
    "  Args:\n",
    "    working_dir (string) - directory to place TF Transform outputs\n",
    "  \n",
    "  Returns:\n",
    "    transform_fn - transformation graph\n",
    "    transformed_train_data - transformed training examples\n",
    "    transformed_test_data - transformed test examples\n",
    "    transformed_metadata - transform output metadata\n",
    "  '''\n",
    "\n",
    "  # Delete TF Transform if it already exists\n",
    "  if os.path.exists(working_dir):\n",
    "    shutil.rmtree(working_dir)\n",
    "\n",
    "  with beam.Pipeline() as pipeline:\n",
    "      with tft_beam.Context(temp_dir=os.path.join(working_dir, TRANSFORM_TEMP_DIR)):\n",
    "        \n",
    "        # Read the input CSV and clean the data\n",
    "        raw_data = (\n",
    "              pipeline\n",
    "              | 'ReadTrainData' >> beam.io.ReadFromText(INPUT_FILE, coder=beam.coders.BytesCoder(), skip_header_lines=1)\n",
    "              | 'CleanLines' >> beam.Map(clean_fn))\n",
    "\n",
    "        # Partition the dataset into train and test sets using the partition_fn defined earlier.    \n",
    "        raw_train_data, raw_test_data = (raw_data\n",
    "                                 | 'TrainTestSplit' >> beam.Partition(partition_fn, 2))\n",
    "        \n",
    "        # Create a TFXIO to read the data with the schema. You need\n",
    "        # to list all columns in order since the schema doesn't specify the\n",
    "        # order of columns in the csv.\n",
    "        csv_tfxio = tfxio.BeamRecordCsvTFXIO(\n",
    "              physical_format='text',\n",
    "              column_names=ordered_columns,\n",
    "              schema=RAW_DATA_SCHEMA)\n",
    "\n",
    "        # Parse the raw train data into inputs for TF Transform\n",
    "        raw_train_data = (raw_train_data \n",
    "                          | 'DecodeTrainData' >> csv_tfxio.BeamSource())\n",
    "        \n",
    "        # Get the raw data metadata\n",
    "        RAW_DATA_METADATA = csv_tfxio.TensorAdapterConfig()\n",
    "        \n",
    "        # Pair the train data with the metadata into a tuple\n",
    "        raw_train_dataset = (raw_train_data, RAW_DATA_METADATA)\n",
    "\n",
    "        # Training data transformation. The TFXIO (RecordBatch) output format\n",
    "        # is chosen for improved performance.\n",
    "        (transformed_train_data,transformed_metadata) , transform_fn = (\n",
    "        raw_train_dataset | tft_beam.AnalyzeAndTransformDataset(preprocessing_fn, output_record_batches=True))\n",
    "\n",
    "\n",
    "        # Parse the raw data into inputs for TF Transform\n",
    "        raw_test_data = (raw_test_data\n",
    "                         | 'DecodeTestData' >> csv_tfxio.BeamSource())\n",
    "        \n",
    "        # Pair the test data with the metadata into a tuple\n",
    "        raw_test_dataset = (raw_test_data, RAW_DATA_METADATA)\n",
    "        \n",
    "        # Now apply the same transform function to the test data.\n",
    "        # You don't need the transformed data schema. It's the same as before.\n",
    "        transformed_test_data, _ = (\n",
    "          (raw_test_dataset, transform_fn) | tft_beam.TransformDataset(output_record_batches=True))\n",
    "        \n",
    "        # Declare an encoder to convert output record batches to TF Examples \n",
    "        transformed_data_coder = RecordBatchToExamplesEncoder(transformed_metadata.schema)\n",
    "        \n",
    "        # Encode transformed train data and write to disk\n",
    "        _ = (\n",
    "            transformed_train_data\n",
    "            | 'EncodeTrainData' >> beam.FlatMapTuple(lambda batch, _: transformed_data_coder.encode(batch))\n",
    "            | 'WriteTrainData' >> beam.io.WriteToTFRecord(\n",
    "                os.path.join(working_dir, TRANSFORM_TRAIN_FILENAME)))\n",
    "\n",
    "        # Encode transformed test data and write to disk\n",
    "        _ = (\n",
    "            transformed_test_data\n",
    "            | 'EncodeTestData' >> beam.FlatMapTuple(lambda batch, _: transformed_data_coder.encode(batch))\n",
    "            | 'WriteTestData' >> beam.io.WriteToTFRecord(\n",
    "                os.path.join(working_dir, TRANSFORM_TEST_FILENAME)))\n",
    "        \n",
    "        # Write transform function to disk\n",
    "        _ = (\n",
    "          transform_fn\n",
    "          | 'WriteTransformFn' >>\n",
    "          tft_beam.WriteTransformFn(os.path.join(working_dir)))\n",
    "\n",
    "         \n",
    "  return transform_fn, transformed_train_data, transformed_test_data, transformed_metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8GPbh3qKzyTb"
   },
   "outputs": [],
   "source": [
    "def main():\n",
    "  return read_and_transform_data(WORKING_DIR)\n",
    "\n",
    "if __name__ == '__main__':\n",
    "  transform_fn, transformed_train_data, trainsformed_test_data, transformed_metadata = main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-CqjN6qrTzQ5"
   },
   "source": [
    "##Prepare Training and Test Datasets from TFTransformOutput\n",
    "\n",
    "Now that you have the transformed dataset, you will need to map it to training and test datasets whch can be used for training a model using TensorFlow. Since this is time series data, it makes sense to group a fixed-length series of measurements and map that to the label found in a future time step. For example, 3 days of data can be used to predict the next day's humidity. You can use the [tf.data.Dataset.window()](https://www.tensorflow.org/api_docs/python/tf/data/Dataset#window) method to implement these groupings.\n",
    "\n",
    "In this exercise, you will use data from the last 5 days to predict the temperature 12 hours into the future. \n",
    "\n",
    "* You already know that data is taken every 10 minutes so there will be 720 data points for 5 days. \n",
    "* You will also assume that drastic change is not expected within an hour so you will downsample the 720 data points to 120 by just getting the data every hour. Thus, you will use a *stride* (or step size) of 6 when getting the data points. This makes the `HISTORY_SIZE` for the feature window equal to 120. \n",
    "* For this single step prediction model (forecasting), the label for a dataset window will be the temperature 12 hours into the future.\n",
    "* By default, the `window()` method moves `size` (i.e. window size) steps at a time. For example, if you have a dataset with elements `[1, 2, 3, 4]` and you have `size=2`, then your results will look like: `[1, 2], [3, 4]`. You can reconfigure this behavior and move at smaller or larger steps with the `shift` parameter. For the same sample dataset with window `size=2` but with `shift=1`, the results will look like `[1, 2], [2,3], [3,4]` because the window is moving `1` element at a time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ho06omK_OAVH"
   },
   "outputs": [],
   "source": [
    "# Constants to prepare the transformed data for modeling\n",
    "\n",
    "LABEL_KEY = 'T (degC)'\n",
    "OBSERVATIONS_PER_HOUR = 6\n",
    "HISTORY_SIZE = 120\n",
    "FUTURE_TARGET = 12\n",
    "BATCH_SIZE = 72\n",
    "SHIFT = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9J8ljQwSQh8b"
   },
   "outputs": [],
   "source": [
    "# Get the output of the Transform component\n",
    "tf_transform_output = tft.TFTransformOutput(os.path.join(WORKING_DIR))\n",
    "\n",
    "# Get the index of the label key\n",
    "index_of_label = list(tf_transform_output.transformed_feature_spec().keys()).index(LABEL_KEY)\n",
    "print(index_of_label)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "S_qnkrUJt9wO"
   },
   "source": [
    "Next, you will define several helper functions to extract the data from your transformed dataset and group it into windows. First, this `parse_function()` will help in getting the transformed features and rearranging them so that the label values (i.e. `T (degC)`) is at the end of the tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XQxYFAXkNUwP"
   },
   "outputs": [],
   "source": [
    "def parse_function(example_proto):\n",
    "    \n",
    "    feature_spec = tf_transform_output.transformed_feature_spec()\n",
    "    \n",
    "    # Define features with the example_proto (transformed data) and the feature_spec using tf.io.parse_single_example \n",
    "    features = tf.io.parse_single_example(example_proto, feature_spec)\n",
    "    values = list(features.values())\n",
    "    values[index_of_label], values[len(features) - 1] = values[len(features) - 1], values[index_of_label]\n",
    "    \n",
    "    # Stack the values along the first axis\n",
    "    stacked_features = tf.stack(values, axis=0)\n",
    "\n",
    "    return stacked_features\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BNkCb1gkPQ0R"
   },
   "source": [
    "Next, you will separate the features and target values into a tuple with this utility function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "T3xcDvgIOrvl"
   },
   "outputs": [],
   "source": [
    "def map_features_target(elements):\n",
    "    features = elements[:HISTORY_SIZE]\n",
    "    target = elements[-1:,-1]\n",
    "    return (features, target)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "abOkQXarPIYl"
   },
   "source": [
    "Finally, you can define the dataset window with the function below. It uses the parameters defined above and also the helper functions to produce the batched feature-target mappings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dAUbyWGfPUha"
   },
   "outputs": [],
   "source": [
    "def get_windowed_dataset(path):\n",
    "        \n",
    "    # Instantiate a tf.data.TFRecordDataset passing in the appropiate path\n",
    "    dataset = tf.data.TFRecordDataset(path)\n",
    "    \n",
    "    # Use the dataset's map method to map the parse_function\n",
    "    dataset = dataset.map(parse_function)\n",
    "    \n",
    "    # Use the window method with expected total size. Define stride and set drop_remainder to True\n",
    "    dataset = dataset.window(HISTORY_SIZE + FUTURE_TARGET, shift=SHIFT, stride=OBSERVATIONS_PER_HOUR, drop_remainder=True)\n",
    "    \n",
    "    # Use the flat_map method passing in an anonymous function that given a window returns window.batch(HISTORY_SIZE + FUTURE_TARGET)\n",
    "    dataset = dataset.flat_map(lambda window: window.batch(HISTORY_SIZE + FUTURE_TARGET))\n",
    "    \n",
    "    # Use the map method passing in the previously defined map_features_target function\n",
    "    dataset = dataset.map(map_features_target) \n",
    "    \n",
    "    # Use the batch method and pass in the appropiate batch size\n",
    "    dataset = dataset.batch(BATCH_SIZE)\n",
    "\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "U-Afrxg2QYeT"
   },
   "source": [
    "You can now use the `get_dataset()` function on your transformed examples to produce the dataset windows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "i0grFNuaTsk3"
   },
   "outputs": [],
   "source": [
    "# Get list of train and test data tfrecord filenames from the transform outputs\n",
    "train_tfrecord_files = tf.io.gfile.glob(os.path.join(WORKING_DIR, TRANSFORM_TRAIN_FILENAME + '*'))\n",
    "test_tfrecord_files = tf.io.gfile.glob(os.path.join(WORKING_DIR, TRANSFORM_TEST_FILENAME + '*'))\n",
    "\n",
    "# Generate dataset windows\n",
    "windowed_train_dataset = get_windowed_dataset(train_tfrecord_files[0])\n",
    "windowed_test_dataset = get_windowed_dataset(test_tfrecord_files[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1AmBzD-OQjMG"
   },
   "source": [
    "Let's preview the resulting shapes of the data windows. If you print the shapes of the tensors in a single batch, you'll notice that it indeed produced the required dimensions. It has 72 examples per batch where each contain 120 measurements for each of the 13 features in the transformed dataset. The target tensor shape only has one feature per example in the batch as expected (i.e. only `T (degC)`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tgTJDfGpt8af"
   },
   "outputs": [],
   "source": [
    "ordered_feature_spec_names = tf_transform_output.transformed_feature_spec().keys()\n",
    "\n",
    "# Preview an example in the train dataset\n",
    "for features, target  in windowed_train_dataset.take(1):\n",
    "    print(f'Shape of input features for a batch: {features.shape}')\n",
    "    print(f'Shape of targets for a batch: {target.shape}\\n')\n",
    "\n",
    "    print(f'INPUT FEATURES:')\n",
    "    for value, name in zip(features[0][0].numpy(), ordered_feature_spec_names):\n",
    "      print(f'{name} : {value}') \n",
    "  \n",
    "    print(f'\\nTARGET TEMPERATURE: {target[0][0]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "60TEQ9tcMORG"
   },
   "outputs": [],
   "source": [
    "# Preview an example in the test dataset\n",
    "for features, target in windowed_test_dataset.take(1):\n",
    "    print(f'Shape of input features for a batch: {features.shape}')\n",
    "    print(f'Shape of targets for a batch: {target.shape}\\n')\n",
    "\n",
    "    print(f'INPUT FEATURES:')\n",
    "    for value, name in zip(features[0][0].numpy(), ordered_feature_spec_names):\n",
    "      print(f'{name} : {value}') \n",
    "  \n",
    "    print(f'\\nTARGET TEMPERATURE: {target[0][0]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wsqC0wqmk9h4"
   },
   "source": [
    "## Wrap Up\n",
    "\n",
    "In this notebook, you got to see how you may want to prepare seasonal data. It shows how you can handle periodicity and produce batches of dataset windows. You will be doing something similar in the next lab with sensor data. This time though, the measurements are taken at a much higher rate (20 Hz). The labels are also categorical so you will be handling that differently.\n",
    "\n",
    "On to the next!"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "C2_W4_Lab_1_WeatherData.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
